import gradio as gr
from sklearn.cluster import KMeans
from PIL import Image
import numpy as np

def extract_palette(image_array, k=5):
    # 确保图像数据是uint8类型
    image_array = image_array.astype(np.uint8)
    # 重塑数组以适应KMeans算法
    data = image_array.reshape((-1, 3))
    # 创建KMeans模型并拟合数据
    kmeans = KMeans(n_clusters=k, n_init=10)
    kmeans.fit(data)
    # 获取调色板
    palette = kmeans.cluster_centers_.astype(int)
    return palette

def process_image(image_array):
    palette = extract_palette(image_array, k=5)
    # 将调色板转换为十六进制颜色代码
    hex_codes = ["#{:02x}{:02x}{:02x}".format(color[0], color[1], color[2]) for color in palette]
    # 创建颜色样本的画廊
    color_samples = [Image.new('RGB', (100, 100), tuple(palette[i])) for i in range(palette.shape[0])]
    # 将PIL图像对象转换为Gradio可以识别的格式
    images_for_gallery = [np.array(img) for img in color_samples]
    return "\n".join(hex_codes), images_for_gallery

# 创建Gradio接口
iface = gr.Interface(
    fn=process_image,
    inputs=gr.Image(label="Upload an image"),  # 移除了shape参数
    outputs=[gr.Text(label="色彩十六进制码"), gr.Gallery(label="调色板预览")],
    title="KMeans调色板提取器",
    description="上传图像获取其调色板."
)

if __name__ == "__main__":
    iface.launch()
